﻿#pragma once

#include <iostream>
#include <cassert>
using namespace std;

namespace TL
{ 
	template<class k, class v>
	struct AVLTreeNode
	{
		pair<k, v> _kv;
		AVLTreeNode<k,v>* _parent;
		AVLTreeNode<k,v>* _left;
		AVLTreeNode<k, v>* _right;
		int _bf;

		AVLTreeNode(const pair<k,v>& kv)
			:_kv(kv)
			,_parent(nullptr)
			,_left(nullptr)
			,_right(nullptr)
			,_bf(0)
		{}
	};

	template<class k, class v>
	class AVLTree
	{
		typedef AVLTreeNode<k, v> Node;
	public:
		//右单旋
		void RotateR(Node* parent)
		{
			Node* ppNode = parent->_parent;
 			Node* subL = parent->_left;
			Node* subLR = subL->_right;

			parent->_left = subLR;
			subL->_right = parent;
			if (parent == _root)
			{
				_root = subL;
			}
			else
			{
				if (ppNode->_left == parent)
					ppNode->_left = subL;
				else
					ppNode->_right = subL;
			}
			parent->_parent = subL;
			subL->_parent = ppNode;
			if (subLR)
				subLR->_parent = parent;

			subL->_bf = parent->_bf = 0;
		}

		//左单旋
		void RotateL(Node* parent)
		{
			Node* ppNode = parent->_parent;
			Node* subR = parent->_right;
			Node* subRL = subR->_left;

			parent->_right = subRL;
			subR->_left = parent;
			
			if (parent == _root)
			{
				_root = subR;
			}
			else
			{
				if (ppNode->_left == parent)
					ppNode->_left = subR;
				else
					ppNode->_right = subR;
			}
			parent->_parent = subR;
			subR->_parent = ppNode;
			if (subRL)
				subRL->_parent = parent;

			subR->_bf = parent->_bf = 0;
		}

		//左右双旋
		void RotateLR(Node* parent)
		{
			Node* subL = parent->_left;
			Node* subLR = subL->_right;
			int bf = subLR->_bf;

			RotateL(subL);
			RotateR(parent);

			if (bf == 0)
			{
				parent->_bf = 0;
				subL->_bf = 0;
				subLR->_bf = 0;
			}
			else if (bf == -1)
			{
				subL->_bf = 0;
				parent->_bf = 1;
				subLR->_bf = 0;
			}
			else if (bf == 1)
			{
				subL->_bf = -1;
				subLR->_bf = 0;
				parent->_bf = 0;
			}
			else
			{
				assert(false);
			}
		}

		//右左双旋
		void RotateRL(Node* parent)
		{
			Node* subR = parent->_right;
			Node* subRL = subR->_left;
			//这里要记录subRL的平衡因子
			//因为在旋转过程中会修改它
			int bf = subRL->_bf;

			RotateR(subR);
			RotateL(parent);

			if (bf == 0)
			{
				parent->_bf = 0;
				subR->_bf = 0;
				subRL->_bf = 0;
			}
			else if (bf == -1)
			{
				subRL->_bf = 0;
				subR->_bf = 1;
				parent->_bf = 0;
			}
			else if (bf == 1)
			{
				subRL->_bf = 0;
				subR->_bf = 0;
				parent->_bf = -1;
			}
			else
			{
				assert(false);
			}
		}


		bool insert(const pair<k, v>& kv)
		{
			Node* cur = _root;
			Node* parent = _root;
			while (cur)
			{
				if (kv.first > cur->_kv.first)
				{
					parent = cur;
					cur = cur->_right;
				}
				else if (kv.first < cur->_kv.first)
				{
					parent = cur;
					cur = cur->_left;
				}
				else
				{
					return false;
				}
			}
			cur = new Node(kv);
			if(parent == nullptr)
			{
				_root = cur;
			}
			else
			{
				cur->_parent = parent;
				if (kv.first > parent->_kv.first)
				{
					parent->_right = cur;
				}
				else
				{
					parent->_left = cur;
				}

				while (cur != _root)
				{
					if (cur == parent->_left)
					{
						parent->_bf--;
					}
					else
					{
						parent->_bf++;
					}

					if (parent->_bf == 0)
					{
						break;
					}
					else if (parent->_bf == -1 || parent->_bf == 1)
					{
						cur = parent;
						parent = cur->_parent;
					}
					else if (parent->_bf == -2 || parent->_bf == 2)
					{
						//旋转调整
						if (parent->_bf == -2 && cur->_bf == -1)
						{
							//纯粹左边高(右单旋)
							RotateR(parent);
						}
						else if (parent->_bf == 2 && cur->_bf == 1)
						{
							//纯粹右边高(左单旋)
							RotateL(parent);
						}
						else if (parent->_bf == -2 && cur->_bf == 1)
						{
							//左边的右边高(左右双旋)
							RotateLR(parent);
						}
						else if (parent->_bf == 2 && cur->_bf == -1)
						{
							//右边的左边高(右左双旋)
							RotateRL(parent);
						}
						else
						{
							assert(false);
						}
						break;
					}
					else
					{
						assert(false);
					}
				}
			}
			return true;
		}

		Node* find(const k& x)
		{
			Node* cur = _root;
			while (cur)
			{
				if (x < cur->_kv.first)
				{
					cur = cur->_left;
				}
				else if(x > cur->_kv.first)
				{
					cur = cur->_right;
				}
				else
				{
					return cur;
				}
			}
			return nullptr;
		}

		void InOrder()
		{
			_InOrder(_root);
		}

		int size()
		{
			return _size(_root);
		}

		int height()
		{
			return _height(_root);
		}

		bool IsBalance()
		{
			return _IsBalance(_root);
		}

	private:

		void _InOrder(Node* root)
		{
			if (root == nullptr)
			{
				return;
			}
			_InOrder(root->_left);
			cout << root->_kv.first << " : " << root->_kv.second << endl;
			//cout << "平衡因子 : " << root->_bf << endl;
			_InOrder(root->_right);
		}

		int _size(Node* root)
		{
			if (root == nullptr)
			{
				return 0;
			}
			return _size(root->_left) + _size(root->_right) + 1;
		}

		int _height(Node* root)
		{
			if (root == nullptr)
			{
				return 0;
			}
			int leftheight = _height(root->_left);
			int rightheight = _height(root->_right);
			return leftheight > rightheight ? leftheight + 1 : rightheight + 1;
		}

		bool _IsBalance(Node* root)
		{
			if (root == nullptr)
			{
				return true;
			}
			int left = _height(root->_left);
			int right = _height(root->_right);
			int realbf = right - left;
			if (abs(realbf) > 1)
			{
				cout << "高度差异常" << endl;
				return false;
			}
			if (root->_bf != realbf)
			{
				cout << root->_kv.first << "平衡因子异常" << endl;
				return false;
			}
			return _IsBalance(root->_left) && _IsBalance(root->_right);
		}

		Node* _root = nullptr;
	};
}